from dataclasses import dataclass  # 导入dataclass装饰器，用于简化创建主要存储数据的类

from langgraph.graph.state import CompiledStateGraph   # 从LangGraph导入编译后的状态图类型
from langgraph.pregel import Pregel   # 从LangGraph导入Pregel图类型

# 导入一系列预定义的各种功能的智能体图
from agents.bg_task_agent.bg_task_agent import bg_task_agent  # 后台任务智能体
from agents.chatbot import chatbot  # 简单聊天机器人智能体
from agents.command_agent import command_agent  # 命令处理智能体
from agents.interrupt_agent import interrupt_agent  # 中断处理智能体
from agents.knowledge_base_agent import kb_agent  # 基于知识库的智能体
from agents.langgraph_supervisor_agent import langgraph_supervisor_agent    # LangGraph监督者智能体
from agents.langgraph_supervisor_hierarchy_agent import langgraph_supervisor_hierarchy_agent   # 层级式LangGraph监督者智能体
from agents.rag_assistant import rag_assistant    # 检索增强生成(RAG)智能体
from agents.research_assistant import research_assistant  # 研究助手智能体（具备网络搜索和计算器等功能）
from agents.time_get import time_get_agent  # 研究助手智能体（具备网络搜索和计算器等功能）
from agents.main_agent import master_agent  # 研究助手智能体（具备网络搜索和计算器等功能）
from schema import AgentInfo   # 从本地schema模块导入AgentInfo数据模型

DEFAULT_AGENT = "research-assistant"

# Type alias to handle LangGraph's different agent patterns
# - @entrypoint functions return Pregel
# - StateGraph().compile() returns CompiledStateGraph
AgentGraph = CompiledStateGraph | Pregel
# 定义类型别名 `AgentGraph`，表示一个智能体图可以是 CompiledStateGraph 或 Pregel 类型
# 这统一了LangGraph中两种创建图的方式（使用@entrypoint装饰器返回Pregel，使用StateGraph().compile()返回CompiledStateGraph

@dataclass
class Agent:
    description: str
    graph: AgentGraph


agents: dict[str, Agent] = {
    "chatbot": Agent(description="A simple chatbot.", graph=chatbot),
    "research-assistant": Agent(
        description="A research assistant with web search and calculator.", graph=research_assistant
    ),
    "rag-assistant": Agent(
        description="A RAG assistant with access to information in a database.", graph=rag_assistant
    ),
    "command-agent": Agent(description="A command agent.", graph=command_agent),
    "bg-task-agent": Agent(description="A background task agent.", graph=bg_task_agent),
    "langgraph-supervisor-agent": Agent(
        description="A langgraph supervisor agent", graph=langgraph_supervisor_agent
    ),
    "langgraph-supervisor-hierarchy-agent": Agent(
        description="A langgraph supervisor agent with a nested hierarchy of agents",
        graph=langgraph_supervisor_hierarchy_agent,
    ),
    "interrupt-agent": Agent(description="An agent the uses interrupts.", graph=interrupt_agent),
    "knowledge-base-agent": Agent(
        description="A retrieval-augmented generation agent using Amazon Bedrock Knowledge Base",
        graph=kb_agent,
    ),
    "time-get-agent": Agent(
        description="An agent specialized in retrieving and providing accurate current time information, including full datetime, year-month-day, day of the week, and handling relative time queries (e.g., 'tomorrow'). It ensures time data is based on the system's current time in the East 8th Time Zone (Beijing Time).",
        graph=time_get_agent,
    ),
    "master-agent": Agent(
        description="A master agent that orchestrates and manages the execution of sub-agents within a multi-agent system. It handles task delegation, coordination, and integration of results from specialized agents to achieve complex objectives efficiently.",
        graph=master_agent,
    ),
}


def get_agent(agent_id: str) -> AgentGraph:
    return agents[agent_id].graph


def get_all_agent_info() -> list[AgentInfo]:
    return [
        AgentInfo(key=agent_id, description=agent.description) for agent_id, agent in agents.items()
    ]
